import jieba  # 分词
import sqlite3
from matplotlib import pyplot as plt  # 绘图，数据可视化
from wordcloud import WordCloud  # 词云
from PIL import Image  # 图片处理
import numpy as np  # 矩阵运算
from collections import Counter  # 用于词频统计


def get_movie_introductions():
    """从SQLite数据库获取电影介绍文本"""
    con = sqlite3.connect('../影片Top.db')
    cur = con.cursor()
    sql = 'SELECT introduction FROM movie'
    data = cur.execute(sql)
    text = ' '.join(item[0] for item in data if item[0])
    cur.close()
    con.close()
    return text


def generate_wordcloud(text):
    """分词并过滤停用词"""
    words = jieba.cut(text)
    stopwords = {'的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着', '看', '好', '这'}
    filtered_words = [word for word in words if word not in stopwords and len(word) > 1]
    return filtered_words


def save_word_frequency_to_db(words, db_path='../word_frequency.db'):
    """统计词频并保存到SQLite数据库"""
    word_count = Counter(words)
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS word_frequency
                 (word TEXT PRIMARY KEY, 
                 count INTEGER)''')
    for word, count in word_count.items():
        c.execute("INSERT OR REPLACE INTO word_frequency (word, count) VALUES (?, ?)", (word, count))
    conn.commit()
    conn.close()
    print(f"已成功将 {len(word_count)} 个词的词频数据保存到 {db_path}")
    return word_count


def create_wordcloud_image(words, mask_image_path='../static/img/词云背景图.jpg', output_path='../static/img/word.jpg'):
    """生成词云图片"""
    text = ' '.join(words)

    try:
        img = Image.open(mask_image_path)
        img_array = np.array(img)
    except Exception as e:
        print(f"无法加载背景图片: {e}，将使用默认样式")
        img_array = None

    wc = WordCloud(
        background_color='white',
        mask=img_array,
        font_path='msyh.ttc',
        width=800,
        height=600,
        max_font_size=150
    )

    wc.generate(text)
    plt.figure(figsize=(10, 8))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"词云图片已保存到 {output_path}")
    plt.show()


def main():
    text = get_movie_introductions()
    words = generate_wordcloud(text)
    word_count = save_word_frequency_to_db(words)
    create_wordcloud_image(words)
    print(f"总词数: {len(words)}")
    print(f"不重复词数: {len(word_count)}")
    print("最常见的10个词:")
    for word, count in word_count.most_common(10):
        print(f"{word}: {count}次")


if __name__ == "__main__":
    main()